/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import org.openhitls.core.CryptoNative;

import javax.net.ssl.*;
import java.nio.ByteBuffer;
import java.nio.ReadOnlyBufferException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.io.ByteArrayInputStream;
import java.util.List;
import java.util.ArrayList;

/**
 * JSSE SSLEngine implementation backed by openHiTLS.
 *
 * <p>This class provides non-blocking TLS operations through the standard
 * Java SSLEngine API. It uses memory BIO operations to interface with
 * the native openHiTLS library.</p>
 *
 * <h2>Usage Pattern</h2>
 * <pre>
 * SSLEngine engine = sslContext.createSSLEngine();
 * engine.setUseClientMode(true);
 * engine.beginHandshake();
 *
 * // Handshake loop
 * while (engine.getHandshakeStatus() != HandshakeStatus.FINISHED) {
 *     SSLEngineResult result = engine.wrap(appOut, netOut);
 *     // Send netOut to network...
 *
 *     result = engine.unwrap(netIn, appIn);
 *     // Process received data...
 * }
 * </pre>
 */
public class HiTlsSSLEngine extends SSLEngine {

    // Parent context
    private final HiTlsSSLContextSpi context;

    // Native SSL session pointer
    private long nativeSslPtr;

    // Connection state
    private boolean clientMode = true;
    private boolean handshakeStarted = false;
    private boolean handshakeCompleted = false;
    private boolean isInboundDone = false;
    private boolean isOutboundDone = false;
    private boolean needClientAuth = false;
    private boolean wantClientAuth = false;

    // Session
    private HiTlsSSLSession session;

    // Configuration
    private String[] enabledCipherSuites;
    private String[] enabledProtocols;
    private SSLParameters sslParameters;

    // Internal buffers for BIO operations
    private static final int DEFAULT_BUFFER_SIZE = 16384;
    private byte[] netInBuffer;
    private byte[] netOutBuffer;
    private int netInBufferPos;
    private int netOutBufferLen;

    // Handshake status
    private SSLEngineResult.HandshakeStatus handshakeStatus = SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;

    /**
     * Creates a new SSLEngine.
     *
     * @param context Parent SSL context
     * @param peerHost Peer hostname (for SNI)
     * @param peerPort Peer port
     */
    HiTlsSSLEngine(HiTlsSSLContextSpi context, String peerHost, int peerPort) {
        super(peerHost, peerPort);
        this.context = context;
        this.enabledCipherSuites = context.getEnabledCipherSuites();
        this.enabledProtocols = context.getEnabledProtocols();
        this.sslParameters = new SSLParameters();
        this.netInBuffer = new byte[DEFAULT_BUFFER_SIZE];
        this.netOutBuffer = new byte[DEFAULT_BUFFER_SIZE];
        this.netInBufferPos = 0;
        this.netOutBufferLen = 0;
    }

    // ==================== Core SSLEngine Operations ====================

    @Override
    public SSLEngineResult wrap(ByteBuffer[] srcs, int offset, int length, ByteBuffer dst)
            throws SSLException {
        checkNotClosed();

        if (srcs == null || dst == null) {
            throw new IllegalArgumentException("Buffers cannot be null");
        }
        if (dst.isReadOnly()) {
            throw new ReadOnlyBufferException();
        }

        // Ensure SSL session is initialized
        ensureInitialized();

        int bytesConsumed = 0;
        int bytesProduced = 0;
        SSLEngineResult.Status status = SSLEngineResult.Status.OK;

        // Handle handshake if not complete
        if (!handshakeCompleted) {
            // Perform handshake step
            int hsResult = CryptoNative.sslDoHandshake(nativeSslPtr);

            if (hsResult == CryptoNative.TLS_SUCCESS) {
                handshakeCompleted = true;
                handshakeStatus = SSLEngineResult.HandshakeStatus.FINISHED;
                updateSession();
            } else if (hsResult == CryptoNative.TLS_WANT_READ) {
                handshakeStatus = SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
            } else if (hsResult == CryptoNative.TLS_WANT_WRITE) {
                handshakeStatus = SSLEngineResult.HandshakeStatus.NEED_WRAP;
            } else {
                throw new SSLException("Handshake failed: " + CryptoNative.sslGetErrorString(nativeSslPtr));
            }
        }

        // If handshake complete, encrypt application data
        if (handshakeCompleted && !isOutboundDone) {
            for (int i = offset; i < offset + length && status == SSLEngineResult.Status.OK; i++) {
                ByteBuffer src = srcs[i];
                if (src.hasRemaining()) {
                    int srcRemaining = src.remaining();
                    byte[] plaintext = new byte[srcRemaining];
                    src.get(plaintext);

                    int written = CryptoNative.sslWrite(nativeSslPtr, plaintext, 0, plaintext.length);
                    if (written > 0) {
                        bytesConsumed += written;
                    } else if (written == CryptoNative.TLS_ERROR) {
                        throw new SSLException("SSL write error");
                    }
                }
            }
        }

        // Read encrypted data from BIO
        int pending = CryptoNative.sslBioPendingWrite(nativeSslPtr);
        if (pending > 0 && dst.remaining() >= pending) {
            byte[] encrypted = new byte[pending];
            int read = CryptoNative.sslBioReadNetData(nativeSslPtr, encrypted, 0, pending);
            if (read > 0) {
                dst.put(encrypted, 0, read);
                bytesProduced = read;
            }
        } else if (pending > 0 && dst.remaining() < pending) {
            status = SSLEngineResult.Status.BUFFER_OVERFLOW;
        }

        // Determine handshake status
        if (handshakeCompleted) {
            handshakeStatus = SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
        }

        return new SSLEngineResult(status, handshakeStatus, bytesConsumed, bytesProduced);
    }

    @Override
    public SSLEngineResult unwrap(ByteBuffer src, ByteBuffer[] dsts, int offset, int length)
            throws SSLException {
        checkNotClosed();

        if (src == null || dsts == null) {
            throw new IllegalArgumentException("Buffers cannot be null");
        }

        // Ensure SSL session is initialized
        ensureInitialized();

        int bytesConsumed = 0;
        int bytesProduced = 0;
        SSLEngineResult.Status status = SSLEngineResult.Status.OK;

        // Feed encrypted data to BIO
        if (src.hasRemaining()) {
            int srcRemaining = src.remaining();
            byte[] encrypted = new byte[srcRemaining];
            src.get(encrypted);

            int written = CryptoNative.sslBioWriteNetData(nativeSslPtr, encrypted, 0, encrypted.length);
            if (written > 0) {
                bytesConsumed = written;
            }
        }

        // Handle handshake if not complete
        if (!handshakeCompleted) {
            int hsResult = CryptoNative.sslDoHandshake(nativeSslPtr);

            if (hsResult == CryptoNative.TLS_SUCCESS) {
                handshakeCompleted = true;
                handshakeStatus = SSLEngineResult.HandshakeStatus.FINISHED;
                updateSession();
            } else if (hsResult == CryptoNative.TLS_WANT_READ) {
                // Check if there's pending output to send first
                int pending = CryptoNative.sslBioPendingWrite(nativeSslPtr);
                if (pending > 0) {
                    // There's data to send - need to wrap first
                    handshakeStatus = SSLEngineResult.HandshakeStatus.NEED_WRAP;
                } else {
                    handshakeStatus = SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
                }
            } else if (hsResult == CryptoNative.TLS_WANT_WRITE) {
                handshakeStatus = SSLEngineResult.HandshakeStatus.NEED_WRAP;
            } else if (hsResult == CryptoNative.TLS_ERROR) {
                throw new SSLException("Handshake failed: " + CryptoNative.sslGetErrorString(nativeSslPtr));
            }
        }

        // If handshake complete, decrypt application data
        if (handshakeCompleted && !isInboundDone) {
            for (int i = offset; i < offset + length && status == SSLEngineResult.Status.OK; i++) {
                ByteBuffer dst = dsts[i];
                if (dst.isReadOnly()) {
                    throw new ReadOnlyBufferException();
                }

                if (dst.hasRemaining()) {
                    byte[] plaintext = new byte[dst.remaining()];
                    int read = CryptoNative.sslRead(nativeSslPtr, plaintext, 0, plaintext.length);

                    if (read > 0) {
                        dst.put(plaintext, 0, read);
                        bytesProduced += read;
                    } else if (read == CryptoNative.TLS_CLOSED) {
                        isInboundDone = true;
                        status = SSLEngineResult.Status.CLOSED;
                        break;
                    } else if (read == CryptoNative.TLS_WANT_READ) {
                        // Need more data
                        status = SSLEngineResult.Status.BUFFER_UNDERFLOW;
                        break;
                    }
                }
            }
        }

        // Determine handshake status
        if (handshakeCompleted) {
            handshakeStatus = SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING;
        }

        return new SSLEngineResult(status, handshakeStatus, bytesConsumed, bytesProduced);
    }

    // ==================== Handshake Management ====================

    @Override
    public void beginHandshake() throws SSLException {
        checkNotClosed();

        if (handshakeStarted && handshakeCompleted) {
            // Renegotiation - not typically supported in TLS 1.3
            throw new SSLException("Renegotiation is not supported");
        }

        ensureInitialized();
        handshakeStarted = true;
        handshakeStatus = clientMode ? SSLEngineResult.HandshakeStatus.NEED_WRAP : SSLEngineResult.HandshakeStatus.NEED_UNWRAP;
    }

    @Override
    public SSLEngineResult.HandshakeStatus getHandshakeStatus() {
        return handshakeStatus;
    }

    // ==================== Session Management ====================

    @Override
    public SSLSession getSession() {
        if (session == null) {
            // Return a dummy session before handshake completes
            session = new HiTlsSSLSession(this);
        }
        return session;
    }

    @Override
    public SSLSession getHandshakeSession() {
        if (!handshakeCompleted) {
            return getSession();
        }
        return null;
    }

    private void updateSession() {
        if (session == null) {
            session = new HiTlsSSLSession(this);
        }
        session.updateFromNative(nativeSslPtr);
    }

    // ==================== Configuration ====================

    @Override
    public void setUseClientMode(boolean mode) {
        if (handshakeStarted) {
            throw new IllegalArgumentException("Cannot change mode after handshake started");
        }
        this.clientMode = mode;
    }

    @Override
    public boolean getUseClientMode() {
        return clientMode;
    }

    @Override
    public void setNeedClientAuth(boolean need) {
        this.needClientAuth = need;
        if (nativeSslPtr != 0) {
            CryptoNative.sslCtxSetVerifyMode(context.getNativeContext(), need ? 2 : 0);
        }
    }

    @Override
    public boolean getNeedClientAuth() {
        return needClientAuth;
    }

    @Override
    public void setWantClientAuth(boolean want) {
        this.wantClientAuth = want;
        if (nativeSslPtr != 0) {
            CryptoNative.sslCtxSetVerifyMode(context.getNativeContext(), want ? 1 : 0);
        }
    }

    @Override
    public boolean getWantClientAuth() {
        return wantClientAuth;
    }

    @Override
    public void setEnableSessionCreation(boolean flag) {
        // Session creation control - typically always enabled
    }

    @Override
    public boolean getEnableSessionCreation() {
        return true;
    }

    @Override
    public String[] getEnabledCipherSuites() {
        return enabledCipherSuites.clone();
    }

    @Override
    public void setEnabledCipherSuites(String[] suites) {
        this.enabledCipherSuites = suites.clone();
    }

    @Override
    public String[] getSupportedCipherSuites() {
        return context.getSupportedCipherSuites();
    }

    @Override
    public String[] getEnabledProtocols() {
        return enabledProtocols.clone();
    }

    @Override
    public void setEnabledProtocols(String[] protocols) {
        this.enabledProtocols = protocols.clone();
    }

    @Override
    public String[] getSupportedProtocols() {
        return context.getSupportedProtocols();
    }

    @Override
    public SSLParameters getSSLParameters() {
        SSLParameters params = new SSLParameters();
        params.setCipherSuites(enabledCipherSuites);
        params.setProtocols(enabledProtocols);
        params.setNeedClientAuth(needClientAuth);
        params.setWantClientAuth(wantClientAuth);
        if (getPeerHost() != null) {
            params.setServerNames(List.of(new SNIHostName(getPeerHost())));
        }
        return params;
    }

    @Override
    public void setSSLParameters(SSLParameters params) {
        if (params.getCipherSuites() != null) {
            setEnabledCipherSuites(params.getCipherSuites());
        }
        if (params.getProtocols() != null) {
            setEnabledProtocols(params.getProtocols());
        }
        setNeedClientAuth(params.getNeedClientAuth());
        setWantClientAuth(params.getWantClientAuth());
        this.sslParameters = params;
    }

    // ==================== Close Operations ====================

    @Override
    public void closeInbound() throws SSLException {
        isInboundDone = true;
    }

    @Override
    public boolean isInboundDone() {
        return isInboundDone;
    }

    @Override
    public void closeOutbound() {
        if (!isOutboundDone && nativeSslPtr != 0) {
            CryptoNative.sslShutdown(nativeSslPtr);
        }
        isOutboundDone = true;
    }

    @Override
    public boolean isOutboundDone() {
        return isOutboundDone;
    }

    @Override
    public Runnable getDelegatedTask() {
        // openHiTLS doesn't require delegated tasks
        return null;
    }

    // ==================== Internal Methods ====================

    private void ensureInitialized() throws SSLException {
        if (nativeSslPtr == 0) {
            nativeSslPtr = CryptoNative.sslNew(context.getNativeContext());
            if (nativeSslPtr == 0) {
                throw new SSLException("Failed to create SSL session");
            }

            // Set client/server mode
            CryptoNative.sslSetConnectState(nativeSslPtr,
                clientMode ? CryptoNative.TLS_MODE_CLIENT : CryptoNative.TLS_MODE_SERVER);

            // Set SNI hostname for client mode
            if (clientMode && getPeerHost() != null) {
                CryptoNative.sslSetHostname(nativeSslPtr, getPeerHost());
            }
        }
    }

    private void checkNotClosed() throws SSLException {
        if (isInboundDone && isOutboundDone) {
            throw new SSLException("SSLEngine is closed");
        }
    }

    /**
     * Get the native SSL pointer.
     * Package-private for use by SSLSession.
     */
    long getNativeSslPtr() {
        return nativeSslPtr;
    }

    /**
     * Get the negotiated key exchange group ID.
     * @return group ID, or 0 if not available
     */
    public int getNegotiatedGroup() {
        if (nativeSslPtr == 0) {
            return 0;
        }
        return CryptoNative.sslGetNegotiatedGroup(nativeSslPtr);
    }

    /**
     * Get the negotiated key exchange group name.
     * @return group name string
     */
    public String getNegotiatedGroupName() {
        int group = getNegotiatedGroup();
        switch (group) {
            case CryptoNative.GROUP_SECP256R1: return "secp256r1";
            case CryptoNative.GROUP_SECP384R1: return "secp384r1";
            case CryptoNative.GROUP_SECP521R1: return "secp521r1";
            case CryptoNative.GROUP_X25519: return "x25519";
            case CryptoNative.GROUP_X448: return "x448";
            case CryptoNative.GROUP_SM2: return "SM2";
            case CryptoNative.GROUP_X25519_MLKEM768: return "X25519_MLKEM768 (Post-Quantum Hybrid)";
            case CryptoNative.GROUP_ECDH_NISTP256_MLKEM768: return "ECDH_NISTP256_MLKEM768 (Post-Quantum Hybrid)";
            case CryptoNative.GROUP_ECDH_NISTP384_MLKEM1024: return "ECDH_NISTP384_MLKEM1024 (Post-Quantum Hybrid)";
            default: return "Unknown (" + group + ")";
        }
    }

    /**
     * Get peer certificates from the native SSL session.
     */
    X509Certificate[] getPeerCertificates() throws SSLException {
        if (nativeSslPtr == 0) {
            return null;
        }

        byte[][] certsData = CryptoNative.sslGetPeerCertificateChain(nativeSslPtr);
        if (certsData == null || certsData.length == 0) {
            return null;
        }

        try {
            CertificateFactory cf = CertificateFactory.getInstance("X.509");
            List<X509Certificate> certs = new ArrayList<>();

            for (byte[] certData : certsData) {
                X509Certificate cert = (X509Certificate) cf.generateCertificate(
                    new ByteArrayInputStream(certData));
                certs.add(cert);
            }

            return certs.toArray(new X509Certificate[0]);
        } catch (Exception e) {
            throw new SSLException("Failed to parse peer certificates", e);
        }
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            if (nativeSslPtr != 0) {
                CryptoNative.sslFree(nativeSslPtr);
                nativeSslPtr = 0;
            }
        } finally {
            super.finalize();
        }
    }
}
